"""
硅基流动（SiliconFlow）图片生成器
支持 FLUX.1-schnell 等模型，每天 14400 次生成限制
"""

import os
import requests
import base64
import time
from typing import Optional, Callable
from datetime import datetime


class SiliconFlowGenerator:
    """硅基流动图片生成器"""
    
    BASE_URL = "https://api.siliconflow.cn/v1/images/generations"
    
    # 支持的模型列表
    MODELS = [
        "Kwai-Kolors/Kolors",              # 快手可图，中文友好，免费 400 张/天
        "black-forest-labs/FLUX.1-schnell",  # FLUX 快速版，免费 400 张/天
        "black-forest-labs/FLUX.1-dev",      # FLUX 开发版，收费
    ]

    # 免费模型（用于限流时优先在免费模型之间切换）
    FREE_MODELS = [
        "Kwai-Kolors/Kolors",
        "black-forest-labs/FLUX.1-schnell",
    ]
    
    def __init__(self, api_key, save_folder: str = "./siliconflow_images", proxies: Optional[dict] = None):
        """
        初始化生成器
        
        Args:
            api_key: SiliconFlow API Key 或 API Keys 列表（支持多个轮换）
            save_folder: 图片保存根目录路径（内部会创建 completed/ 与 tmp/ 子目录）
            proxies: 可选代理配置，例如 {"http": "socks5h://127.0.0.1:1086", "https": "socks5h://127.0.0.1:1086"}
        """
        # 支持单个 API Key 或多个 API Keys
        if isinstance(api_key, list):
            self.api_keys = api_key
            self.current_key_index = 0
        else:
            self.api_keys = [api_key]
            self.current_key_index = 0
        
        self.save_folder = save_folder  # 根目录
        self.proxies = proxies  # 调用 API / 下载图片时使用的代理（默认直连）
        self._ensure_folder_exists()
    
    def _get_current_api_key(self) -> str:
        """获取当前使用的 API Key"""
        return self.api_keys[self.current_key_index]
    
    def _rotate_api_key(self):
        """轮换到下一个 API Key"""
        if len(self.api_keys) > 1:
            self.current_key_index = (self.current_key_index + 1) % len(self.api_keys)
    
    def _ensure_folder_exists(self):
        """确保保存文件夹及子目录存在"""
        os.makedirs(self.save_folder, exist_ok=True)
        os.makedirs(os.path.join(self.save_folder, "completed"), exist_ok=True)
        os.makedirs(os.path.join(self.save_folder, "tmp"), exist_ok=True)

    def _get_alternative_model(self, current_model: str) -> Optional[str]:
        """在免费模型之间切换：

        - 当前模型是 FREE_MODELS 之一，则返回列表中的下一个模型
        - 否则返回 None（例如付费模型时不自动改模型）
        """
        try:
            if current_model not in self.FREE_MODELS:
                return None
            idx = self.FREE_MODELS.index(current_model)
            # 只有两个免费模型时，这里就是在两者之间互换
            next_idx = (idx + 1) % len(self.FREE_MODELS)
            if next_idx == idx:
                return None
            return self.FREE_MODELS[next_idx]
        except Exception:
            return None
    
    def generate_image(
        self,
        prompt: str,
        width: int = 576,
        height: int = 1024,
        model: str = "black-forest-labs/FLUX.1-schnell",
        num_inference_steps: int = 20,
        guidance_scale: float = 7.5,
        progress_callback: Optional[Callable[[str], None]] = None
    ) -> Optional[str]:
        """
        生成单张图片
        
        Args:
            prompt: 图片描述
            width: 图片宽度
            height: 图片高度
            model: 模型名称
            num_inference_steps: 推理步数
            guidance_scale: 引导比例
            progress_callback: 进度回调函数
            
        Returns:
            保存的图片路径，失败返回 None
        """
        try:
            if progress_callback:
                progress_callback(f"正在生成图片: {prompt[:50]}...")
            
            # 构建请求（使用当前 API Key）
            current_api_key = self._get_current_api_key()
            headers = {
                "Authorization": f"Bearer {current_api_key}",
                "Content-Type": "application/json"
            }
            
            payload = {
                "model": model,
                "prompt": prompt,
                "image_size": f"{width}x{height}",
                "num_inference_steps": num_inference_steps,
                "guidance_scale": guidance_scale,
                "batch_size": 1
            }
            
            # 智能退避重试：对 429 限流错误进行重试
            # 策略：
            #   1) 先在免费模型之间切换（Kolors ⇄ FLUX.1-schnell）
            #   2) 仍然限流时，再轮换 API Key
            #   3) 之后按随机退避时间重试
            max_backoff_retries = 3
            backoff_retry_count = 0
            switched_model_once = False  # 每次调用最多自动切换一次模型
            
            response = None
            result = None
            
            while backoff_retry_count <= max_backoff_retries:
                # 内层：超时重试循环
                max_timeout_retries = 2
                for timeout_attempt in range(max_timeout_retries + 1):
                    try:
                        response = requests.post(
                            self.BASE_URL,
                            headers=headers,
                            json=payload,
                            timeout=120,
                            proxies=self.proxies,
                        )
                        break  # 请求成功，跳出超时重试循环
                    except (requests.exceptions.Timeout, requests.exceptions.ConnectionError) as e:
                        if timeout_attempt >= max_timeout_retries:
                            if progress_callback:
                                progress_callback("超时，已跳过")
                            return None
                        if progress_callback:
                            progress_callback(f"超时，{timeout_attempt + 1}/{max_timeout_retries + 1} 重试...")
                        time.sleep(2)
                
                # 如果没有获取到响应，返回 None
                if response is None:
                    return None
                
                # 检查响应状态码
                if response.status_code == 429:
                    # 限流错误：优先在免费模型之间切换，然后再轮换 API Key
                    # 1) 先尝试切换模型（仅在免费模型之间，并且每次调用最多切一次）
                    if not switched_model_once:
                        alternative = self._get_alternative_model(payload["model"])
                        if alternative:
                            old_model = payload["model"]
                            payload["model"] = alternative
                            switched_model_once = True
                            if progress_callback:
                                progress_callback(f"限流，先切换模型 {old_model} → {alternative}")
                            # 短暂停顿后重新请求（使用同一个 API Key）
                            time.sleep(1)
                            continue
                    
                    # 2) 仍然限流：第一次先轮换 API Key
                    if backoff_retry_count == 0 and len(self.api_keys) > 1:
                        self._rotate_api_key()
                        current_api_key = self._get_current_api_key()
                        headers["Authorization"] = f"Bearer {current_api_key}"
                        if progress_callback:
                            progress_callback(f"限流，切换到第 {self.current_key_index + 1} 个 API Key")
                        backoff_retry_count += 1
                        time.sleep(2)  # 短暂停顿
                        continue
                    
                    # 3) 继续限流：进入随机退避重试
                    if backoff_retry_count < max_backoff_retries:
                        import random
                        wait_time = random.randint(20, 40)  # 20-40 秒随机
                        if progress_callback:
                            progress_callback(f"限流，等待{wait_time}秒")
                        time.sleep(wait_time)
                        backoff_retry_count += 1
                        continue
                    else:
                        if progress_callback:
                            progress_callback("限流，已跳过")
                        return None
                
                elif response.status_code == 502 or response.status_code == 503:
                    # 服务器临时错误 - 使用较短的退避时间
                    if backoff_retry_count < max_backoff_retries:
                        wait_time = 5 * (2 ** backoff_retry_count)  # 5秒, 10秒, 20秒
                        if progress_callback:
                            progress_callback(f"{response.status_code}错误，等待{wait_time}秒")
                        time.sleep(wait_time)
                        backoff_retry_count += 1
                        continue  # 继续下一次退避重试
                    else:
                        if progress_callback:
                            progress_callback(f"{response.status_code}错误，已跳过")
                        return None
                
                # 其他错误正常抛出，由外层处理
                response.raise_for_status()
                result = response.json()
                
                # 成功获取响应，跳出重试循环
                break
            
            # 提取图片数据
            if "images" in result and len(result["images"]) > 0:
                image_data = result["images"][0]
                
                # 判断是 URL 还是 base64
                if "url" in image_data:
                    # 下载图片
                    img_response = requests.get(
                        image_data["url"], timeout=60, proxies=self.proxies
                    )
                    img_response.raise_for_status()
                    image_bytes = img_response.content
                elif "b64_json" in image_data:
                    # 解码 base64
                    image_bytes = base64.b64decode(image_data["b64_json"])
                else:
                    raise Exception("无法获取图片数据")
                
                # 生成文件名（使甤13位毫秒时间戳，与样例格式一致）
                timestamp_ms = int(time.time() * 1000)  # 13位毫秒时间戳
                filename = f"{timestamp_ms}.jpeg"
                completed_dir = os.path.join(self.save_folder, "completed")
                tmp_dir = os.path.join(self.save_folder, "tmp")
                os.makedirs(completed_dir, exist_ok=True)
                os.makedirs(tmp_dir, exist_ok=True)
                final_path = os.path.join(completed_dir, filename)
                tmp_path = os.path.join(tmp_dir, filename + ".part")
                
                # 保存图片为 JPEG 格式（先写入临时文件，验证成功后再原子重命名）
                from PIL import Image
                from io import BytesIO
                try:
                    # 先用 verify 检查图片是否完整、格式正确
                    img_for_verify = Image.open(BytesIO(image_bytes))
                    img_for_verify.verify()
                    # verify 之后需要重新打开一次
                    img = Image.open(BytesIO(image_bytes))
                    # 转换为 RGB（如果是 RGBA）
                    if img.mode in ("RGBA", "LA", "P"):
                        background = Image.new("RGB", img.size, (255, 255, 255))
                        if img.mode == "P":
                            img = img.convert("RGBA")
                        background.paste(img, mask=img.split()[-1] if img.mode in ("RGBA", "LA") else None)
                        img = background
                    elif img.mode != "RGB":
                        img = img.convert("RGB")
                    # 保存到临时文件
                    img.save(tmp_path, "JPEG", quality=95, optimize=True)
                    # 原子替换到最终路径，避免出现半张图片
                    os.replace(tmp_path, final_path)
                except ImportError:
                    # 如果没有 PIL，只能保证字节写入完整，同样采用临时文件 + 原子重命名
                    with open(tmp_path, "wb") as f:
                        f.write(image_bytes)
                    os.replace(tmp_path, final_path)
                except Exception:
                    # 任意异常都不保留半成品
                    try:
                        if os.path.exists(tmp_path):
                            os.remove(tmp_path)
                    except Exception:
                        pass
                    raise
                
                if progress_callback:
                    progress_callback(f"✓ 图片已保存: {filename}")
                
                return final_path
            else:
                raise Exception("API 返回数据格式错误")
                
        except requests.exceptions.HTTPError as e:
            return None
        except Exception as e:
            return None
    
    def generate_batch(
        self,
        prompts: list[str],
        width: int = 576,
        height: int = 1024,
        model: str = "black-forest-labs/FLUX.1-schnell",
        images_per_prompt: int = 1,
        num_inference_steps: int = 20,
        guidance_scale: float = 7.5,
        progress_callback: Optional[Callable[[str], None]] = None,
        should_stop: Optional[Callable[[], bool]] = None
    ) -> list[str]:
        """
        批量生成图片
        
        Args:
            prompts: 关键词列表
            width: 图片宽度
            height: 图片高度
            model: 模型名称
            images_per_prompt: 每个关键词生成的图片数量
            num_inference_steps: 推理步数
            guidance_scale: 引导比例
            progress_callback: 进度回调函数
            should_stop: 停止检查函数，如果返回 True 则停止生成
            
        Returns:
            成功保存的图片路径列表
        """
        saved_images = []
        total = len(prompts) * images_per_prompt
        current = 0
        
        if progress_callback:
            progress_callback(f"开始批量生成，共 {total} 张图片")
        
        for prompt in prompts:
            # 检查是否需要停止
            if should_stop and should_stop():
                if progress_callback:
                    progress_callback("已停止生成")
                break
            
            for i in range(images_per_prompt):
                # 检查是否需要停止
                if should_stop and should_stop():
                    if progress_callback:
                        progress_callback("已停止生成")
                    break
                
                current += 1
                
                if progress_callback:
                    progress_callback(f"进度: {current}/{total} - {prompt[:30]}...")
                
                # 生成图片
                filepath = self.generate_image(
                    prompt=prompt,
                    width=width,
                    height=height,
                    model=model,
                    num_inference_steps=num_inference_steps,
                    guidance_scale=guidance_scale,
                    progress_callback=None  # 内部不再回调，避免重复
                )
                
                if filepath:
                    saved_images.append(filepath)
                    if progress_callback:
                        saved_count = len(saved_images)
                        progress_callback(f"✓ [{saved_count}/{total}] 已保存")
                else:
                    if progress_callback:
                        progress_callback(f"✗ [{current}/{total}] 失败")
        
        if progress_callback:
            progress_callback(f"批量生成完成！成功: {len(saved_images)}/{total}")
        
        return saved_images
    
    @staticmethod
    def get_available_models() -> list[str]:
        """获取可用的模型列表"""
        return SiliconFlowGenerator.MODELS.copy()
    
    @staticmethod
    def validate_api_key(api_key: str) -> bool:
        """
        验证 API Key 是否有效
        
        Args:
            api_key: API Key
            
        Returns:
            是否有效
        """
        try:
            headers = {
                "Authorization": f"Bearer {api_key}",
                "Content-Type": "application/json"
            }
            
            # 发送一个简单的测试请求
            response = requests.get(
                "https://api.siliconflow.cn/v1/models",
                headers=headers,
                timeout=10
            )
            
            return response.status_code == 200
        except:
            return False
